import math
import heapq
from tqdm import tqdm
import numpy as np
import pickle

train_path = "../data/train.txt"
test_path = "../data/test.txt"
sim_path = "sim.pkl"
result_path = "result.txt"

start_line = 0 
end_line = 0
user_item_list = {}
item_set = set()

with open(train_path, 'r') as file: 
    train_lines = file.readlines()

for index, line in enumerate(train_lines):
    if '|' in line:
        userid, num_rating_item = line.split('|')
        start_line = index + 1
        end_line = index + int(num_rating_item)
        user_item_list[userid] = {}
    elif index >= start_line and index <= end_line:
        line = line.strip('\n')
        item, rating = line.split("  ")
        user_item_list[userid][item] = int(rating)
        item_set.add(item)

print("Finished loading train data")

user_len = len(user_item_list)
print("User Num:", user_len)
print("Rating Num:", len(train_lines) - len(user_item_list))
print("Item Num:", len(item_set))

item_user_list = {}
for userid, item_list in user_item_list.items():
    for item in item_list:
        item_user_list.setdefault(item, dict())
        item_user_list[item][userid] = item_list[item]

item_rating_avg = {}
for item, user_list in item_user_list.items():
    avg = 0.0
    for user, rating_score in user_list.items():
        avg += rating_score
    avg = avg / len(user_list)
    item_rating_avg[item] = avg
    
print("Finished calculating item rating avg")

user_rating_avg = {}
all_avg = 0.0
for userid, item_list in user_item_list.items():
    avg = 0.0
    for item, rating_score in item_list.items():
        avg += rating_score
        all_avg += rating_score
    avg = avg / len(item_list)
    user_rating_avg[userid] = avg
all_avg = all_avg / (len(train_lines) - len(user_item_list))

print(all_avg)
print(user_rating_avg["0"])
print("Finished calculating user rating avg and all rating avg")

user_norm = {}
for userid, item_list in user_item_list.items():
    sum = 0.0
    for item, rating_score in item_list.items():
        rating_score = rating_score - user_rating_avg[userid]
        user_item_list[userid][item] = rating_score
        sum += rating_score * rating_score
    user_norm[userid] = math.sqrt(sum)

print(user_norm["0"])
print(user_item_list["0"]["518385"])
print("Finished calculating user norm")

similarity = {}
for i, (userid1, item_list1) in tqdm(enumerate(user_item_list.items()), total=user_len):
    similarity[userid1] = {}
    for j, (userid2, item_list2) in enumerate(user_item_list.items()):
        if user_norm[userid1] * user_norm[userid2] == 0:
            similarity[userid1][userid2] = 0.0
            continue
        if i > j:
            similarity[userid1][userid2] = similarity[userid2][userid1]
        elif i == j:
            similarity[userid1][userid2] = 1.0
        else:
            cos_sim = 0.0
            for item, rating in item_list1.items():
                if item in item_list2:
                    cos_sim += rating  * item_list2[item] 
            cos_sim = cos_sim / (user_norm[userid1] * user_norm[userid2])
            similarity[userid1][userid2] = cos_sim


with open(test_path, 'r') as file: 
    test_lines = file.readlines()

start_line = 0 
end_line = 0
test_user_item_list = {}
    
for index, line in enumerate(test_lines):
    if '|' in line:
        userid, num_rating_item = line.split('|')
        start_line = index + 1
        end_line = index + int(num_rating_item)
        test_user_item_list[userid] = {}
    elif index >= start_line and index <= end_line:
        line = line.strip('\n')
        test_user_item_list[userid][line] = int(-1)

for userid, item_list in tqdm(test_user_item_list.items(), total=len(test_user_item_list)):
    for item in item_list:
        largest_n = heapq.nlargest(10, ((k, v) for k, v in similarity[userid].items() if item in user_item_list[k]), key=lambda item: item[1])
        rating_sum = 0.0
        sim_sum = 0.0
        for user_sim, value in largest_n:
            rating_sum = rating_sum + similarity[userid][user_sim] * user_item_list[user_sim][item]
            sim_sum = sim_sum + similarity[userid][user_sim]
        if sim_sum == 0:
            test_user_item_list[userid][item] = user_rating_avg[userid]
        else:
            test_user_item_list[userid][item] = (rating_sum / sim_sum) + user_rating_avg[userid]
        if test_user_item_list[userid][item] < 0:
            test_user_item_list[userid][item] = user_rating_avg[userid]
        if test_user_item_list[userid][item] > 100:
            test_user_item_list[userid][item] = 100

count = 0
with open(result_path, 'w') as f:
    for userid, item_list in test_user_item_list.items():
        for item, rating in item_list.items():
            if count == 0:
                f.write(f"{userid}|6\n")
            f.write(f"{item}  {rating}\n")
            count = count + 1
            count = count % 6


